#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function

# todo 采集数据
import cv2
import sql
import numpy as np
import time
from robot import r2t,t2r,transl
from common import draw_str,drawChessboard

def chaiji(nao_name):
    w,h = 12,9
    unit = 20.0
    #w,h = 9,4
    #unit = 15.0
    nao = sql.Nao.get_by_name(nao_name)
    autonomousLifeProxy = nao._get_proxy("ALAutonomousLife")
    postureProxy = nao._get_proxy("ALRobotPosture")
    step = np.deg2rad(3)
    minHeadYaw, maxHeadYaw, minHeadPitch,maxHeadPitch = None,None,None,None
    while True:
        top_img = nao.top_cam.get_image()
        bottom_img = nao.bottom_cam.get_image()
        motionProxy = nao._get_proxy("ALMotion")
        HeadYaw,HeadPitch = motionProxy.getAngles("Head",True)
        Tb_o,Tc_o,position = nao.findTr(w,h,unit,top_img,bottom_img)
        drawChessboard(top_img,w=12,h=9)
        drawChessboard(bottom_img,w=12,h=9)
        draw_str(top_img,(10,10),"HeadYaw:"+str(HeadYaw))
        draw_str(top_img,(10,25),"HeadPitch:"+str(HeadPitch))
        draw_str(top_img,(10,40),"minHeadYaw:"+str(minHeadYaw))
        draw_str(top_img,(10,55),"maxHeadYaw:"+str(maxHeadYaw))
        draw_str(top_img,(10,70),"minHeadPitch:"+str(minHeadPitch))
        draw_str(top_img,(10,85),"maxHeadPitch:"+str(maxHeadPitch))
        draw_str(bottom_img,(10,10),"Tb_o:")
        if Tb_o is not None:
            for x,p in enumerate(np.round(Tb_o,3).tolist()):
                draw_str(bottom_img,(10,25+x*15),str(p))
        draw_str(bottom_img,(10,90),"Tc_o:")
        if Tc_o is not None:
            for x,p in enumerate(np.round(Tc_o,3).tolist()):
                draw_str(bottom_img,(10,105+x*15),str(p))

        cv2.imshow('top', top_img)
        cv2.imshow('bottom', bottom_img)
        setAngle = lambda anglse: motionProxy.setAngles('Head', anglse, 0.5)
        ch = cv2.waitKey(100) & 0xFF

        command = {
            "q": lambda step:setAngle([HeadYaw+step,HeadPitch]),
            "a": lambda step:setAngle([HeadYaw-step,HeadPitch]),
            "z": lambda step:setAngle([0,HeadPitch]),
            "w": lambda step:setAngle([HeadYaw,HeadPitch+step]),
            "s": lambda step:setAngle([HeadYaw,HeadPitch-step]),
            "x": lambda step:setAngle([HeadYaw,0]),
            "i": lambda step:autonomousLifeProxy.setState("disabled"),
            "o": lambda step:motionProxy.wakeUp(),
            "p": lambda step:motionProxy.rest(),
            "j": lambda step:postureProxy.goToPosture("Stand", 1.0),
            "k": lambda step:postureProxy.goToPosture("StandInit", 1.0),
            "l": lambda step:postureProxy.goToPosture("Crouch", 1.0),
        }
        if ch == 27:
            break
        elif chr(ch) in command:
            func = command[chr(ch)]
            func(step)
        elif ch == ord('e'): # 设定各角的极限
            maxHeadYaw = HeadYaw
        elif ch == ord('d'):
            minHeadYaw = HeadYaw
        elif ch == ord('r'):
            maxHeadPitch = HeadPitch
        elif ch == ord('f'):
            minHeadPitch = HeadPitch
        elif ch == ord("c"): # 保存当前
            nao.states.append(sql.State("cam",nao,w,h,unit))
            ss.add(nao)
            ss.commit()
            ss = sql.Session()
        elif ch == ord('v'): # 按照极限保存
            ss = sql.Session()
            for headYaw in np.linspace(minHeadYaw,maxHeadYaw,10):
                for headPitch in np.linspace(minHeadPitch,maxHeadPitch,10):
                    print(headYaw,headPitch)
                    setAngle([headYaw,headPitch])
                    time.sleep(10)
                    nao.states.append(sql.State("cam",nao,w,h,unit))
            ss.add(nao)
            ss.commit()

def connect(nao_name):
    nao = sql.Nao.get_by_name(nao_name)
    nao.connect = True
    ss = sql.Session()
    ss.add(nao)
    ss.commit()

def disconnect(nao_name):
    nao = sql.Nao.get_by_name(nao_name)
    nao.connect = False
    ss = sql.Session()
    ss.add(nao)
    ss.commit()

if __name__ == '__main__':
    nao_name = "vred"
    connect(nao_name)
    chaiji(nao_name)
    disconnect(nao_name)